from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.edge.service import Service
from selenium.common.exceptions import TimeoutException, WebDriverException
from selenium.webdriver.common.action_chains import ActionChains
import time
import os
import subprocess
import winreg
import requests
import zipfile
import platform
import pyautogui
import csv
import sys
pyautogui.PAUSE = 0.1
pyautogui.FAILSAFE = True

# ========== 1. 自动下载 EdgeDriver ==========
def get_edge_version():
    try:
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Edge\BLBeacon") as key:
            version, _ = winreg.QueryValueEx(key, "version")
            return version
    except Exception:
        return None

def download_edgedriver(version: str):
    arch = "win64" if platform.machine().endswith("64") else "win32"
    url = f"https://msedgedriver.azureedge.net/{version}/edgedriver_{arch}.zip"
    driver_dir = os.path.join(os.path.expanduser("~"), "auto_edgedriver")
    os.makedirs(driver_dir, exist_ok=True)
    driver_exe = os.path.join(driver_dir, "msedgedriver.exe")
    # 若已存在且版本匹配则直接返回
    if os.path.isfile(driver_exe):
        return driver_exe
    try:
        resp = requests.get(url, stream=True, timeout=30)
        resp.raise_for_status()
        zip_path = os.path.join(driver_dir, "edgedriver.zip")
        with open(zip_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)
        with zipfile.ZipFile(zip_path) as z:
            z.extractall(driver_dir)
        os.remove(zip_path)
        return driver_exe
    except Exception as e:
        raise RuntimeError(f"下载 EdgeDriver 失败: {e}")

def get_driver_path():
    version = get_edge_version()
    if not version:
        raise RuntimeError("无法读取本机 Edge 版本")
    return download_edgedriver(version)

# ========== 2. 清理残留 Edge 进程 ==========
def kill_edge():
    try:
        subprocess.run(["taskkill", "/F", "/IM", "msedge.exe"],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception:
        pass

# ========== 3. 数据持久化 ==========
CSV_FILE = "comments.csv"

def save_comments_to_csv(driver):
    """增量保存当前已加载的评论到 CSV（追加写）"""
    try:
        comment_items = WebDriverWait(driver, 6).until(
            EC.presence_of_all_elements_located((By.CSS_SELECTOR, 'div[data-e2e="comment-item"]'))
        )
    except TimeoutException:
        comment_items = driver.find_elements(By.CSS_SELECTOR, 'div[data-e2e="comment-item"]')

    if not comment_items:
        return 0

    # 读取已有用户名去重
    seen = set()
    if os.path.isfile(CSV_FILE):
        with open(CSV_FILE, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                seen.add(row["username"])

    new_rows = 0
    with open(CSV_FILE, "w", newline="", encoding="utf-8") as f:
        writer = None
        for item in comment_items:
            try:
                username = item.find_element(
                    By.CSS_SELECTOR, 'a[href*="/user/"] span span span span span span span'
                ).text.strip()
            except Exception:
                continue
            if username in seen:
                continue
            seen.add(username)

            try:
                user_link = item.find_element(By.CSS_SELECTOR, 'a[href*="/user/"]').get_attribute("href")
            except Exception:
                user_link = ""

            try:
                time_location = item.find_element(By.CSS_SELECTOR, "div.fJhvAqos span").text
                location = time_location.split("·")[-1].strip()
            except Exception:
                location = ""

            try:
                remark = next(
                    (tag.text.strip() for tag in
                     item.find_elements(By.CSS_SELECTOR, "span.comment-item-tag-text")
                     if tag.text.strip()), "")
            except Exception:
                remark = ""

            row = {
                "username": username,
                "location": location,
                "following_count": "",
                "follower_count": "",
                "follow_btn_status": "unknown",
                "status": f"{username}_未处理",
                "remark": remark,
                "user_link": user_link
            }

            if writer is None:
                writer = csv.DictWriter(f, fieldnames=row.keys())
                if f.tell() == 0:
                    writer.writeheader()
            writer.writerow(row)
            new_rows += 1
    return new_rows

# ========== 4. 主流程 ==========
def main():
    kill_edge()
    user_data_dir = os.path.expanduser(r"~\AppData\Local\Microsoft\Edge\User Data")
    profile_dir = "Default"

    try:
        driver_path = get_driver_path()
        print(f"使用 EdgeDriver: {driver_path}")
    except Exception as e:
        print(f"查找或自动下载失败: {e}")
        # 兜底路径，请按实际手动修改
        driver_path = r"C:\Users\Devin\.cache\selenium\msedgedriver\win64\136.0.3240.47\msedgedriver.exe"
        print(f"使用手改备选驱动: {driver_path}")

    options = webdriver.EdgeOptions()
    options.add_argument(f"--user-data-dir={user_data_dir}")
    options.add_argument(f"--profile-directory={profile_dir}")

    driver = webdriver.Edge(service=Service(driver_path), options=options)
    driver.maximize_window()
    wait = WebDriverWait(driver, 10)

    try:
        driver.get("https://www.douyin.com/user/self")
        # 进入点赞列表（超时则使用pyautogui兜底方案）
        try:
            like_tab = wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, "#semiTablike")))
            like_tab.click()
        except TimeoutException:
            print("未能定位点赞选项卡，启用盲点，请确保元素无遮挡")
            pyautogui.click(618, 460, duration=0.01)

        # 批量获取视频列表并打开第一个
        try:
            user_detail_elements = wait.until(EC.presence_of_all_elements_located(
                (By.CSS_SELECTOR, "div.SfwAcdr1.JrV13Yco")))
            first_video = user_detail_elements[0]
            first_video.click()
        except TimeoutException:
            print("未能定位第一个视频，启用盲点，请确保元素无遮挡")
            pyautogui.click(618, 666, duration=1.01)
        # 点击评论区
        pyautogui.click(2500, 866, duration=1.01)

        # 关闭自动连播
        ActionChains(driver).scroll_by_amount(0, -300).perform()
        autoplay_btn = wait.until(EC.element_to_be_clickable(
            (By.CSS_SELECTOR, "xg-icon.xgplayer-autoplay-setting.automatic-continuous")))
        tip = autoplay_btn.find_element(By.CSS_SELECTOR, "div.xgTips").text.split("\n")[0]
        if tip == "关闭自动连播":
            autoplay_btn.click()
            print("已关闭自动连播")

        # ========== 滚动加载全部评论 ==========
        max_scroll_attempts = 18000
        scroll_count = 0
        loop_count = 0
        try:
            loading_tip = WebDriverWait(driver, 10).until(EC.presence_of_element_located(
                (By.CSS_SELECTOR, "div.ETuXBjRi")))
        except TimeoutException:
            loading_tip = None
            print("未能定位评论加载提示，无法执行滚动加载")
            exit(250)

        while loop_count < max_scroll_attempts:
            loop_count += 1
            print(f"当前第 {loop_count} 次循环，最多循环 {max_scroll_attempts} 次")
            try:
                bottom_tip = driver.find_element(By.CSS_SELECTOR, "div.cnqD49jq")
                print(f"已到达底部: {bottom_tip.text}")
                break
            except TimeoutException:
                driver.execute_script(
                    "arguments[0].scrollIntoView({block: 'nearest', behavior: 'smooth'});",
                    loading_tip)
                scroll_count += 1
                if scroll_count % 1000 == 0:  # 每滚动 1000 次持久化一次
                    rows = save_comments_to_csv(driver)
                    print(f"滚动{scroll_count}次，新增保存{rows}条记录")
                time.sleep(0.3)

        # 最终保存
        rows = save_comments_to_csv(driver)
        print(f"爬取完成，总计保存 {rows} 条新评论")

    except KeyboardInterrupt:
        print("\n用户中断，正在保存已加载数据...")
        rows = save_comments_to_csv(driver)
        print(f"已紧急保存 {rows} 条评论")
    except Exception as e:
        print("发生异常:", e)
        rows = save_comments_to_csv(driver)
        print(f"异常前已保存 {rows} 条评论")
    finally:
        print(f"已保存 {rows} 条评论")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)